package edu.zju.gis.dbfg.model.scala

import edu.zju.gis.dbfg.model.base.GeoItem
import geotrellis.proj4.LatLng
import geotrellis.raster.io.geotiff.GeoTiff
import geotrellis.raster.rasterize.Rasterizer
import geotrellis.raster.rasterize.Rasterizer.foreachCellByGeometry
import geotrellis.raster.{ArrayTile, RasterExtent, Tile, _}
import geotrellis.vector.Extent
import org.apache.spark.Partitioner

import scala.collection.JavaConverters._
import scala.collection.mutable.ArrayBuffer

object MRasterizer{

  def getOptions: Rasterizer.Options = Rasterizer.Options.DEFAULT

  def getPartitioner: Option[Partitioner] = None

  def getRasterExtent(extent: Extent, xresolution: Double, yresolution: Double): RasterExtent = {
    val cols = ((extent.xmax-extent.xmin)/xresolution).toInt + 1
    val rows = ((extent.ymax-extent.ymin)/yresolution).toInt + 1
    RasterExtent(extent, cols, rows)
  }

  def rasterizeWithValue(jFeatures: java.util.List[GeoItem], rasterExtent: RasterExtent, jFields: java.util.List[Tuple2[String,String]]): MultibandTile = {
    val features = jFeatures.asScala
    val fields = jFields.asScala
    val cols = rasterExtent.cols
    var tileBuffer = ArrayBuffer[Tile]()
    for (field <- fields) {
      if (field._2 == "EMPTY") {
        // TODO 不会写 scala 的 continue, 只能这样, do nothing
      } else {
        val feaArray = Array.ofDim[Float](rasterExtent.cols * rasterExtent.rows).fill(floatNODATA)
        for (i <- features.indices) {
          foreachCellByGeometry(features(i).getGeometry, rasterExtent) {
            (col, row) =>
              val data = features(i).getProp.get(field._1)
              if (data == null) {
                feaArray(row * cols + col) = floatNODATA
              } else {
                feaArray(row * cols + col) = data.toString.toFloat
              }
          }
        }
        val tile = ArrayTile(feaArray, rasterExtent.cols, rasterExtent.rows)
        tileBuffer += tile
      }
    }
    ArrayMultibandTile(tileBuffer)
  }

  def write(tile: MultibandTile, extent: Extent, outpath: String) = {
    GeoTiff(tile, extent, LatLng).write(outpath)
  }

}
